import asyncio
import os
import dotenv

from openai import AsyncOpenAI

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()

# Defining multiple agents
History_tutor  = Agent(
    name="History_tutor",
    instructions="""
    You are a history tutor, and your task is to help users solve problems related to historical events, historical figures, and historical dynasties.
You can only answer questions that are directly related to historical events, historical figures, and historical dynasties.
If the user's question is not related to historical events, historical figures, or historical dynasties, please answer directly: I can't do it.
Explain your reasoning process and include examples at each step. Answer in Chinese
    """,
    handoff_description="""
    An agent who specializes in dealing with issues related to historical events, historical figures, historical dynasties, etc.Answer in Chinese
    """
)
Mathematics_tutor = Agent(
    name="Mathematics_tutor",
    handoff_description="Specialized in dealing with mathematical problems agent",
    instructions="You provide help with math problems. Explain your reasoning and include examples at each step.",
)

# Classification Agent
triage_agent = Agent(
    name="Category Agent",
    instructions="""
 You decide which agent to use based on the user's homework question. You can only choose from [History Counselor, Math Counselor], and cannot create a new agent.
You must first output transfer_to_XXX (XXX is the specific agent name), and then output the answer.Answer in Chinese
For example:
- The user asks "How many emperors were there in the Qing Dynasty", you should first output transfer_to_History Counselor, and then output the specific answer.
- The user asks "What is the true meaning of stocks", "Hello translate into English", "Sometimes in life, there must be, the next sentence is", these do not belong to history or mathematics, you should only output 'I can't do it'.
If you can't find the corresponding agent, please only output 'I can't do it'.Answer in Chinese
    """,
    handoffs=[History_tutor, Mathematics_tutor]
)
async def main():
    user_input = input("Please ask your questions:")
    result = await Runner.run(
        starting_agent = triage_agent,
        input =user_input,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    print(result.final_output)

if __name__ == "__main__":  
    asyncio.run(main())